# 以下是一个完整的示例，展示如何使用 BERT 模型获取句子嵌入并计算余弦相似度。
# 为了获取实际的嵌入向量，我们可以使用 Hugging Face 的 transformers 库中的预训练模型。
# 已经提前从HuggingFace下载了bert-base-uncased嵌入模型

from transformers import BertTokenizer, BertModel
import torch
from sklearn.metrics.pairwise import cosine_similarity

# 加载预训练的BERT模型和分词器
tokenizer = BertTokenizer.from_pretrained('./bert-base-uncased')
model = BertModel.from_pretrained('./bert-base-uncased')


def get_embedding(text):
    # 将文本转换为BERT输入格式
    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True)
    # 获取BERT模型输出
    with torch.no_grad():
        outputs = model(**inputs)
    # 取[CLS]标记的嵌入作为句子的表示
    cls_embedding = outputs.last_hidden_state[:, 0, :].numpy()
    return cls_embedding
# 示例句子
sentences = [
    "我喜欢在周末去公园散步。",
    "那里风景优美，空气清新。"
]
# 获取句子嵌入
embeddings = [get_embedding(sentence) for sentence in sentences]
# 计算余弦相似度
similarity = cosine_similarity(embeddings[0], embeddings[1])

print(f"句子 '{sentences[0]}' 和句子 '{sentences[1]}' 的余弦相似度: {similarity[0][0]}")